#!/usr/bin/env python
from __future__ import print_function
import subprocess
import numpy as np
try:
    from astropy.io import fits as pf
except ImportError:
    import pyfits as pf

class EBSimple(object):
    def __init__(self, infile):
        self.infile = infile
        hdulist = pf.open(self.infile)
        try:
            EBoundHDU = hdulist['EBOUNDS']
            self.Emin = np.array([EBoundHDU.data.field(1) * 1.e-3]).flatten() # [MeV]
            self.Emax = np.array([EBoundHDU.data.field(2) * 1.e-3]).flatten() # [MeV]
            self.Energy = np.unique(np.concatenate([self.Emin, self.Emax])) # [MeV]
        except KeyError:
            EBoundHDU = hdulist['ENERGIES']
            self.Energy = EBoundHDU.data.field(0) # [MeV]
            self.Emin = np.array([self.Energy[:-1]]).flatten() # [MeV]
            self.Emax = np.array([self.Energy[1:]]).flatten()  # [MeV]
        hdulist.close()

        self.nEBound = len(self.Emin)
        self.nEnergy = len(self.Energy)


class MapToolsSimple(object):
    def __init__(self, mapfile):
        hdulist = pf.open(mapfile.strip())
        self.map = hdulist[0].data.copy() # map[ilayer, iy, ix]
        self.header = hdulist[0].header

        if self.header['NAXIS'] > 2:
            self.nlayers = self.header['NAXIS3']
        else:
            self.nlayers = 1
            self.map = np.array([self.map])


def mybool(Input):
    return {'True' : True, 'False' : False, 'T' : True, 'F' : False,
            't' : True, 'f' : False, 'TRUE' : True, 'FALSE' : False,
            "true" : True, "false" : False, "1" : True, "0" : False}.get(Input)


def analyseMap(mapcube):
    mmax, mmin = mapcube.max(), mapcube.min()
    mmean, sigma = mapcube.mean(), mapcube.std()
    return mmax, mmin, mmean, sigma


def makeResidue(CCUBE, MCUBE):
    subprocess.call(['farith', CCUBE, MCUBE, 'residue.fits', '-', 'clobber=yes'])
    subprocess.call(['farith', 'residue.fits', MCUBE, 'frac.fits', '/', 'clobber=yes'])


def mapCheck(CCUBE, MCUBE):
    ccube_energy = EBSimple(CCUBE).Energy
    mcube_energy = EBSimple(MCUBE).Energy
    frac_max = (np.abs(ccube_energy-mcube_energy)/ccube_energy).max()
    if frac_max > 1.e-8:
        print('[WARN] The energy in ccube is not compatible with mcube! frac_max = ', frac_max)

    ccube = MapToolsSimple(CCUBE)
    mcube = MapToolsSimple(MCUBE)

    print('-'*100)
    print('[Map Check]')

    chi2 = 0.
    for iE in range(ccube.nlayers):
        print('[%s]: [%s:%s]' % (iE, ccube_energy[iE], ccube_energy[iE+1]))
        ccube_map = ccube.map[iE, :, :]
        mcube_map = mcube.map[iE, :, :]

        ccube_cnts = ccube_map.sum()
        mcube_cnts = mcube_map.sum()
        cnts_frac = (mcube_cnts - ccube_cnts)/ccube_cnts
        chi2 += (ccube_cnts - mcube_cnts)**2
        print('  Model Cnts / Obs Cnts = %i / %i = 1.%+.5f' %  (mcube_cnts, ccube_cnts, cnts_frac))

        resd_map = ccube_map - mcube_map
        frac_map = resd_map / mcube_map
        sigf_map = resd_map / np.sqrt(mcube_map)
        print('  Residue: max=%s, min=%s, mean=%s, sig=%s' % analyseMap(resd_map))
        print('  Fraction: max=%s, min=%s, mean=%s, sig=%s' % analyseMap(frac_map))
        print('  Significance: max=%s, min=%s' % analyseMap(sigf_map)[:2])

    print('-'*100)
    ccube_sum = ccube.map.sum()
    mcube_sum = mcube.map.sum()
    chi2 /= mcube_sum
    print('Model Cnts / Obs Cnts = %i / %i = 1.%+.5f' % (mcube_sum, ccube_sum, (mcube_sum/ccube_sum-1.)))
    print('Chi2 = ', chi2 / (ccube.nlayers-1.))


def specCheck(srcMaps, expCube, binnedExpMap, srcModel, irfs, centerSrc):
    obs = BAn.BinnedObs(srcMaps=srcMaps, expCube=expCube,
                        binnedExpMap=binnedExpMap, irfs=irfs)
    import pyLikelihood as pyLike
    import BinnedAnalysis as BAn
    from LikelihoodState import LikelihoodState

    MIN = BAn.BinnedAnalysis(obs, srcModel, optimizer="MINUIT")
    MINobj = pyLike.Minuit(MIN.logLike)
    sedEnergy = MIN.energies

    if centerSrc in MIN.sourceNames():
        print('[Spectrum Check]')
        TS_fit = MIN.Ts(centerSrc)
        TS_band = 0.

        for src in MIN.sourceNames():
            if MIN[src].src.fixedSpectrum(): continue
            if src == centerSrc:
                if (src_par == 'Prefactor' or src_par == 'norm' or
                    src_par == 'Integral' or src_par == 'Value' or
                    src_par == 'Normalization'):
                    continue
            for src_par in MIN[src].funcs['Spectrum'].paramNames:
                src_index = MIN.par_index(src, src_par)
                MIN.freeze(src_index)

        pristine = LikelihoodState(MIN)
        print('#iband\tFlux\tFluxErr\tTS')
        for iE in range(len(sedEnergy)-1):
            MIN.setEnergyRange(sedEnergy[iE], sedEnergy[iE+1])
            MIN.fit(covar=True, optObject=MINobj, verbosity=0)

            Flux = MIN.flux(centerSrc, emin=sedEnergy[iE], emax=sedEnergy[iE+1], energyFlux=False)
            FluxErr = MIN.fluxError(centerSrc, emin=sedEnergy[iE], emax=sedEnergy[iE+1], energyFlux=False)
            TS = MIN.Ts(centerSrc)
            TS_band += TS
            print('%i\t%s\t%s\t%s' % (iE, Flux, FluxErr, TS))

            pristine.restore()
    else:
        raise RuntimeError('%s is not in the model' % centerSrc)

    print('-'*100)
    print('TS_band - TS_fit = %s - %s = %s' % (TS_band, TS_fit, TS_band-TS_fit))


def cli():
    import argparse

    parser = argparse.ArgumentParser(description='Show the counts in model and ccube')
    parser.add_argument("CCUBE", type=str, help='counts cube')
    parser.add_argument("MCUBE", type=str, help='model counts ccube')
    parser.add_argument("centerSrc", type=str, help='center source')
    parser.add_argument("-makeResidue", type=mybool, default=False, help='Whether to make residue.fits and frac.fits')
    args = parser.parse_args()

    if args.makeResidue:
        makeResidue(args.CCUBE, args.MCUBE)

    mapCheck(args.CCUBE, args.MCUBE)

if __name__ == '__main__': cli()
